// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <string>

#include <pollux/functions/udf_output_string.h>
#include <pollux/common/strings/string_view.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::exec {
    class StringWriter : public UDFOutputString {
    public:
        // Used to initialize top-level strings and allow zero-copy writes.
        StringWriter(FlatVector<StringView> *vector, int32_t offset)
            : vector_(vector), offset_(offset) {
        }

        // If not initialized for zero-copy write, returns a string to copy into the
        // target vector on commit.
        const std::string &value() const {
            return value_;
        }

        /// Reserve a space for the output string with size of at least newCapacity
        void reserve(size_t newCapacity) override {
            if (newCapacity <= capacity()) {
                return;
            }

            auto *newDataBuffer = vector_->getBufferWithSpace(newCapacity);
            auto actualCapacity = newDataBuffer->capacity() - newDataBuffer->size();

            // Impossible to be the same due to the way the capacity is computed.
            DKCHECK(dataBuffer_ != newDataBuffer);

            auto newStartAddress =
                    newDataBuffer->asMutable<char>() + newDataBuffer->size();

            if (size() != 0) {
                std::memcpy(newStartAddress, data(), size());
            }

            setCapacity(actualCapacity);
            setData(newStartAddress);
            dataBuffer_ = newDataBuffer;
        }

        /// Not called by the UDF Implementation. Should be called at the end to
        /// finalize the allocation and the string writing.
        void finalize() {
            if (!finalized_) {
                POLLUX_DCHECK(size() == 0 || data());
                if LIKELY(size()) {
                    DKCHECK(dataBuffer_);
                    dataBuffer_->setSize(dataBuffer_->size() + size());
                }
                vector_->setNoCopy(offset_, StringView(data(), size()));
            }
        }

        void prepare_for_reuse(bool isSet) {
            if (isSet) {
                setCapacity(capacity() - size());
                setData(data() + size());
            }
            resize(0);
            finalized_ = false;
        }

        void finalizeNull() {
        }

        void setEmpty() {
            static const StringView kEmpty("");
            vector_->setNoCopy(offset_, kEmpty);
            finalized_ = true;
        }

        void setNoCopy(const StringView &value) {
            vector_->setNoCopy(offset_, value);
            finalized_ = true;
        }

        template<typename T>
        void operator+=(const T &input) {
            append(input);
        }

        void operator+=(const char *input) {
            append(std::string_view(input));
        }

        template<typename T>
        void operator=(const T &input) {
            resize(0);
            append(input);
        }

        template<typename T>
        void append(const T &input) {
            DKCHECK(!finalized_);
            auto oldSize = size();
            resize(this->size() + input.size());
            if (input.size() != 0) {
                DKCHECK(data());
                DKCHECK(input.data());
                std::memcpy(data() + oldSize, input.data(), input.size());
            }
        }

        void append(const char *input) {
            append(std::string_view(input));
        }

        template<typename T>
        void copy_from(const T &input) {
            append(input);
        }

        void copy_from(const char *input) {
            append(std::string_view(input));
        }

    private:
        StringWriter() = default;

        bool finalized_{false};

        /// The buffer that the output string uses for its allocation set during
        /// reserve() call
        Buffer *dataBuffer_ = nullptr;

        FlatVector<StringView> *vector_;

        int32_t offset_;

        std::string value_;

        template<typename A, typename B>
        friend struct VectorWriter;
    };
} // namespace kumo::pollux::exec
